import os
import pandas as pd
import numpy as np
from PIL import Image
import shutil

def convert_pixels_to_images():
    """将FER2013数据集中的像素数据转换为图像文件"""
    # 加载数据集
    df = pd.read_csv('FERPlus_with_class.csv')
    
    # 基础输出目录
    base_dir = 'data/fer2013'
    
    # 创建文件夹
    for usage in df['Usage'].unique():
        usage_path = os.path.join(base_dir, usage)
        os.makedirs(usage_path, exist_ok=True)
    
    # 转换像素为图像并保存
    for idx, row in df.iterrows():
        pixels = list(map(int, row['pixels'].split()))
        image_array = np.array(pixels).reshape(48, 48).astype(np.uint8)
        image = Image.fromarray(image_array)
        
        usage = row['Usage']
        filename = f"{idx}.png"
        filepath = os.path.join(base_dir, usage, filename)
        image.save(filepath)
    
    print("✅ 转换完成。图像已保存在 data/fer2013/[Usage]/ 文件夹中。")

def organize_images_by_emotion():
    """根据情绪标签组织图像文件"""
    # 基础路径
    base_dir = 'data/fer2013'
    csv_file = 'FERPlus_with_class.csv'
    
    # 情绪标签
    emotion_labels = [
        'neutral', 'happiness', 'surprise', 'sadness',
        'anger', 'disgust', 'fear', 'contempt', 'unknown'
    ]
    
    # 加载新的CSV
    df = pd.read_csv(csv_file)
    
    # 为每个用途和情绪创建子目录
    for usage in df['Usage'].unique():
        if pd.isna(usage): continue
        for emotion in emotion_labels:
            os.makedirs(os.path.join(base_dir, usage, emotion), exist_ok=True)
    
    # 跟踪已移动的文件
    moved_files = set()
    
    # 处理每一行
    for _, row in df.iterrows():
        usage = row['Usage']
        image_name = row['Image name']
        
        if pd.isna(image_name) or row['NF'] == 10:
            continue
        
        # 转换图像名称
        try:
            index = int(image_name.replace('fer', '').replace('.png', ''))
            actual_image_name = f"{index}.png"
        except:
            continue
        
        # 确定主要情绪
        emotion_counts = row[emotion_labels].values
        if sum(emotion_counts) == 0:
            continue
        
        dominant_emotion = emotion_labels[emotion_counts.argmax()]
        
        src_path = os.path.join(base_dir, usage, actual_image_name)
        dst_path = os.path.join(base_dir, usage, dominant_emotion, actual_image_name)
        
        if os.path.exists(src_path):
            shutil.move(src_path, dst_path)
            moved_files.add((usage, actual_image_name))
    
    # 删除未移动的图像
    for usage in ['Training', 'PrivateTest', 'PublicTest']:
        usage_path = os.path.join(base_dir, usage)
        if not os.path.exists(usage_path):
            continue
        
        for file in os.listdir(usage_path):
            file_path = os.path.join(usage_path, file)
            if os.path.isfile(file_path) and (usage, file) not in moved_files:
                os.remove(file_path)
    
    print("✅ 完成：图像已按情绪分类。未标记的文件已删除。")

if __name__ == "__main__":
    # 确保data目录存在
    os.makedirs('data', exist_ok=True)
    
    # 执行数据处理
    convert_pixels_to_images()
    organize_images_by_emotion() 